{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summarization (TensorFlow)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install the Transformers, Datasets, and Evaluate libraries to run this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install datasets evaluate transformers[sentencepiece]\n",
    "!apt install git-lfs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will need to setup git, adapt your email and name in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git config --global user.email \"you@example.com\"\n",
    "!git config --global user.name \"Your Name\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will also need to be logged in to the Hugging Face Hub. Execute the following and enter your credentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['review_id', 'product_id', 'reviewer_id', 'stars', 'review_body', 'review_title', 'language', 'product_category'],\n",
       "        num_rows: 200000\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['review_id', 'product_id', 'reviewer_id', 'stars', 'review_body', 'review_title', 'language', 'product_category'],\n",
       "        num_rows: 5000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['review_id', 'product_id', 'reviewer_id', 'stars', 'review_body', 'review_title', 'language', 'product_category'],\n",
       "        num_rows: 5000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "spanish_dataset = load_dataset(\"amazon_reviews_multi\", \"es\")\n",
    "english_dataset = load_dataset(\"amazon_reviews_multi\", \"en\")\n",
    "english_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'>> Title: Worked in front position, not rear'\n",
       "'>> Review: 3 stars because these are not rear brakes as stated in the item description. At least the mount adapter only worked on the front fork of the bike that I got it for.'\n",
       "\n",
       "'>> Title: meh'\n",
       "'>> Review: Does it’s job and it’s gorgeous but mine is falling apart, I had to basically put it together again with hot glue'\n",
       "\n",
       "'>> Title: Can\\'t beat these for the money'\n",
       "'>> Review: Bought this for handling miscellaneous aircraft parts and hanger \"stuff\" that I needed to organize; it really fit the bill. The unit arrived quickly, was well packaged and arrived intact (always a good sign). There are five wall mounts-- three on the top and two on the bottom. I wanted to mount it on the wall, so all I had to do was to remove the top two layers of plastic drawers, as well as the bottom corner drawers, place it when I wanted and mark it; I then used some of the new plastic screw in wall anchors (the 50 pound variety) and it easily mounted to the wall. Some have remarked that they wanted dividers for the drawers, and that they made those. Good idea. My application was that I needed something that I can see the contents at about eye level, so I wanted the fuller-sized drawers. I also like that these are the new plastic that doesn\\'t get brittle and split like my older plastic drawers did. I like the all-plastic construction. It\\'s heavy duty enough to hold metal parts, but being made of plastic it\\'s not as heavy as a metal frame, so you can easily mount it to the wall and still load it up with heavy stuff, or light stuff. No problem there. For the money, you can\\'t beat it. Best one of these I\\'ve bought to date-- and I\\'ve been using some version of these for over forty years.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def show_samples(dataset, num_samples=3, seed=42):\n",
    "    sample = dataset[\"train\"].shuffle(seed=seed).select(range(num_samples))\n",
    "    for example in sample:\n",
    "        print(f\"\\n'>> Title: {example['review_title']}'\")\n",
    "        print(f\"'>> Review: {example['review_body']}'\")\n",
    "\n",
    "\n",
    "show_samples(english_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "home                      17679\n",
       "apparel                   15951\n",
       "wireless                  15717\n",
       "other                     13418\n",
       "beauty                    12091\n",
       "drugstore                 11730\n",
       "kitchen                   10382\n",
       "toy                        8745\n",
       "sports                     8277\n",
       "automotive                 7506\n",
       "lawn_and_garden            7327\n",
       "home_improvement           7136\n",
       "pet_products               7082\n",
       "digital_ebook_purchase     6749\n",
       "pc                         6401\n",
       "electronics                6186\n",
       "office_product             5521\n",
       "shoes                      5197\n",
       "grocery                    4730\n",
       "book                       3756\n",
       "Name: product_category, dtype: int64"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "english_dataset.set_format(\"pandas\")\n",
    "english_df = english_dataset[\"train\"][:]\n",
    "# Show counts for top 20 products\n",
    "english_df[\"product_category\"].value_counts()[:20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_books(example):\n",
    "    return (\n",
    "        example[\"product_category\"] == \"book\"\n",
    "        or example[\"product_category\"] == \"digital_ebook_purchase\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "english_dataset.reset_format()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'>> Title: I\\'m dissapointed.'\n",
       "'>> Review: I guess I had higher expectations for this book from the reviews. I really thought I\\'d at least like it. The plot idea was great. I loved Ash but, it just didnt go anywhere. Most of the book was about their radio show and talking to callers. I wanted the author to dig deeper so we could really get to know the characters. All we know about Grace is that she is attractive looking, Latino and is kind of a brat. I\\'m dissapointed.'\n",
       "\n",
       "'>> Title: Good art, good price, poor design'\n",
       "'>> Review: I had gotten the DC Vintage calendar the past two years, but it was on backorder forever this year and I saw they had shrunk the dimensions for no good reason. This one has good art choices but the design has the fold going through the picture, so it\\'s less aesthetically pleasing, especially if you want to keep a picture to hang. For the price, a good calendar'\n",
       "\n",
       "'>> Title: Helpful'\n",
       "'>> Review: Nearly all the tips useful and. I consider myself an intermediate to advanced user of OneNote. I would highly recommend.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spanish_books = spanish_dataset.filter(filter_books)\n",
    "english_books = english_dataset.filter(filter_books)\n",
    "show_samples(english_books)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'>> Title: Easy to follow!!!!'\n",
       "'>> Review: I loved The dash diet weight loss Solution. Never hungry. I would recommend this diet. Also the menus are well rounded. Try it. Has lots of the information need thanks.'\n",
       "\n",
       "'>> Title: PARCIALMENTE DAÑADO'\n",
       "'>> Review: Me llegó el día que tocaba, junto a otros libros que pedí, pero la caja llegó en mal estado lo cual dañó las esquinas de los libros porque venían sin protección (forro).'\n",
       "\n",
       "'>> Title: no lo he podido descargar'\n",
       "'>> Review: igual que el anterior'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import concatenate_datasets, DatasetDict\n",
    "\n",
    "books_dataset = DatasetDict()\n",
    "\n",
    "for split in english_books.keys():\n",
    "    books_dataset[split] = concatenate_datasets(\n",
    "        [english_books[split], spanish_books[split]]\n",
    "    )\n",
    "    books_dataset[split] = books_dataset[split].shuffle(seed=42)\n",
    "\n",
    "# Peek at a few examples\n",
    "show_samples(books_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "books_dataset = books_dataset.filter(lambda x: len(x[\"review_title\"].split()) > 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "model_checkpoint = \"google/mt5-small\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [336, 259, 28387, 11807, 287, 62893, 295, 12507, 1], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1]}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs = tokenizer(\"I loved reading the Hunger Games!\")\n",
    "inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['▁I', '▁', 'loved', '▁reading', '▁the', '▁Hung', 'er', '▁Games', '</s>']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.convert_ids_to_tokens(inputs.input_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_input_length = 512\n",
    "max_target_length = 30\n",
    "\n",
    "\n",
    "def preprocess_function(examples):\n",
    "    model_inputs = tokenizer(\n",
    "        examples[\"review_body\"],\n",
    "        max_length=max_input_length,\n",
    "        truncation=True,\n",
    "    )\n",
    "    labels = tokenizer(\n",
    "        examples[\"review_title\"], max_length=max_target_length, truncation=True\n",
    "    )\n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    return model_inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_datasets = books_dataset.map(preprocess_function, batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_summary = \"I absolutely loved reading the Hunger Games\"\n",
    "reference_summary = \"I loved reading the Hunger Games\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install rouge_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "\n",
    "rouge_score = evaluate.load(\"rouge\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'rouge1': AggregateScore(low=Score(precision=0.86, recall=1.0, fmeasure=0.92), mid=Score(precision=0.86, recall=1.0, fmeasure=0.92), high=Score(precision=0.86, recall=1.0, fmeasure=0.92)),\n",
       " 'rouge2': AggregateScore(low=Score(precision=0.67, recall=0.8, fmeasure=0.73), mid=Score(precision=0.67, recall=0.8, fmeasure=0.73), high=Score(precision=0.67, recall=0.8, fmeasure=0.73)),\n",
       " 'rougeL': AggregateScore(low=Score(precision=0.86, recall=1.0, fmeasure=0.92), mid=Score(precision=0.86, recall=1.0, fmeasure=0.92), high=Score(precision=0.86, recall=1.0, fmeasure=0.92)),\n",
       " 'rougeLsum': AggregateScore(low=Score(precision=0.86, recall=1.0, fmeasure=0.92), mid=Score(precision=0.86, recall=1.0, fmeasure=0.92), high=Score(precision=0.86, recall=1.0, fmeasure=0.92))}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scores = rouge_score.compute(\n",
    "    predictions=[generated_summary], references=[reference_summary]\n",
    ")\n",
    "scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Score(precision=0.86, recall=1.0, fmeasure=0.92)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scores[\"rouge1\"].mid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install nltk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "\n",
    "nltk.download(\"punkt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'I grew up reading Koontz, and years ago, I stopped,convinced i had \"outgrown\" him.'\n",
       "'Still,when a friend was looking for something suspenseful too read, I suggested Koontz.'\n",
       "'She found Strangers.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from nltk.tokenize import sent_tokenize\n",
    "\n",
    "\n",
    "def three_sentence_summary(text):\n",
    "    return \"\\n\".join(sent_tokenize(text)[:3])\n",
    "\n",
    "\n",
    "print(three_sentence_summary(books_dataset[\"train\"][1][\"review_body\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_baseline(dataset, metric):\n",
    "    summaries = [three_sentence_summary(text) for text in dataset[\"review_body\"]]\n",
    "    return metric.compute(predictions=summaries, references=dataset[\"review_title\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'rouge1': 16.74, 'rouge2': 8.83, 'rougeL': 15.6, 'rougeLsum': 15.96}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "score = evaluate_baseline(books_dataset[\"validation\"], rouge_score)\n",
    "rouge_names = [\"rouge1\", \"rouge2\", \"rougeL\", \"rougeLsum\"]\n",
    "rouge_dict = dict((rn, round(score[rn].mid.fmeasure * 100, 2)) for rn in rouge_names)\n",
    "rouge_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TFAutoModelForSeq2SeqLM\n",
    "\n",
    "model = TFAutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForSeq2Seq\n",
    "\n",
    "data_collator = DataCollatorForSeq2Seq(tokenizer, model=model, return_tensors=\"tf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_datasets = tokenized_datasets.remove_columns(\n",
    "    books_dataset[\"train\"].column_names\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],\n",
       "        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]), 'input_ids': tensor([[  1494,    259,   8622,    390,    259,    262,   2316,   3435,    955,\n",
       "            772,    281,    772,   1617,    263,    305,  14701,    260,   1385,\n",
       "           3031,    259,  24146,    332,   1037,    259,  43906,    305,    336,\n",
       "            260,      1,      0,      0,      0,      0,      0,      0],\n",
       "        [   259,  27531,  13483,    259,   7505,    260, 112240,  15192,    305,\n",
       "          53198,    276,    259,  74060,    263,    260,    459,  25640,    776,\n",
       "           2119,    336,    259,   2220,    259,  18896,    288,   4906,    288,\n",
       "           1037,   3931,    260,   7083, 101476,   1143,    260,      1]]), 'labels': tensor([[ 7483,   259,  2364, 15695,     1,  -100],\n",
       "        [  259, 27531, 13483,   259,  7505,     1]]), 'decoder_input_ids': tensor([[    0,  7483,   259,  2364, 15695,     1],\n",
       "        [    0,   259, 27531, 13483,   259,  7505]])}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "features = [tokenized_datasets[\"train\"][i] for i in range(2)]\n",
    "data_collator(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_train_dataset = model.prepare_tf_dataset(\n",
    "    tokenized_datasets[\"train\"],\n",
    "    collate_fn=data_collator,\n",
    "    shuffle=True,\n",
    "    batch_size=8,\n",
    ")\n",
    "tf_eval_dataset = model.prepare_tf_dataset(\n",
    "    tokenized_datasets[\"validation\"],\n",
    "    collate_fn=data_collator,\n",
    "    shuffle=False,\n",
    "    batch_size=8,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import create_optimizer\n",
    "import tensorflow as tf\n",
    "\n",
    "# The number of training steps is the number of samples in the dataset, divided by the batch size then multiplied\n",
    "# by the total number of epochs. Note that the tf_train_dataset here is a batched tf.data.Dataset,\n",
    "# not the original Hugging Face Dataset, so its len() is already num_samples // batch_size.\n",
    "num_train_epochs = 8\n",
    "num_train_steps = len(tf_train_dataset) * num_train_epochs\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "\n",
    "optimizer, schedule = create_optimizer(\n",
    "    init_lr=5.6e-5,\n",
    "    num_warmup_steps=0,\n",
    "    num_train_steps=num_train_steps,\n",
    "    weight_decay_rate=0.01,\n",
    ")\n",
    "\n",
    "model.compile(optimizer=optimizer)\n",
    "\n",
    "# Train in mixed-precision float16\n",
    "tf.keras.mixed_precision.set_global_policy(\"mixed_float16\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers.keras_callbacks import PushToHubCallback\n",
    "\n",
    "callback = PushToHubCallback(\n",
    "    output_dir=f\"{model_name}-finetuned-amazon-en-es\", tokenizer=tokenizer\n",
    ")\n",
    "\n",
    "model.fit(\n",
    "    tf_train_dataset, validation_data=tf_eval_dataset, callbacks=[callback], epochs=8\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "generation_data_collator = DataCollatorForSeq2Seq(\n",
    "    tokenizer, model=model, return_tensors=\"tf\", pad_to_multiple_of=320\n",
    ")\n",
    "\n",
    "tf_generate_dataset = model.prepare_tf_dataset(\n",
    "    tokenized_datasets[\"validation\"],\n",
    "    collate_fn=generation_data_collator,\n",
    "    shuffle=False,\n",
    "    batch_size=8,\n",
    "    drop_remainder=True,\n",
    ")\n",
    "\n",
    "\n",
    "@tf.function(jit_compile=True)\n",
    "def generate_with_xla(batch):\n",
    "    return model.generate(\n",
    "        input_ids=batch[\"input_ids\"],\n",
    "        attention_mask=batch[\"attention_mask\"],\n",
    "        max_new_tokens=32,\n",
    "    )\n",
    "\n",
    "\n",
    "all_preds = []\n",
    "all_labels = []\n",
    "for batch, labels in tqdm(tf_generate_dataset):\n",
    "    predictions = generate_with_xla(batch)\n",
    "    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)\n",
    "    labels = labels.numpy()\n",
    "    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "    decoded_preds = [\"\\n\".join(sent_tokenize(pred.strip())) for pred in decoded_preds]\n",
    "    decoded_labels = [\"\\n\".join(sent_tokenize(label.strip())) for label in decoded_labels]\n",
    "    all_preds.extend(decoded_preds)\n",
    "    all_labels.extend(decoded_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = rouge_score.compute(\n",
    "    predictions=decoded_preds, references=decoded_labels, use_stemmer=True\n",
    ")\n",
    "result = {key: value.mid.fmeasure * 100 for key, value in result.items()}\n",
    "{k: round(v, 4) for k, v in result.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "hub_model_id = \"huggingface-course/mt5-small-finetuned-amazon-en-es\"\n",
    "summarizer = pipeline(\"summarization\", model=hub_model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_summary(idx):\n",
    "    review = books_dataset[\"test\"][idx][\"review_body\"]\n",
    "    title = books_dataset[\"test\"][idx][\"review_title\"]\n",
    "    summary = summarizer(books_dataset[\"test\"][idx][\"review_body\"])[0][\"summary_text\"]\n",
    "    print(f\"'>>> Review: {review}'\")\n",
    "    print(f\"\\n'>>> Title: {title}'\")\n",
    "    print(f\"\\n'>>> Summary: {summary}'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'>>> Review: Nothing special at all about this product... the book is too small and stiff and hard to write in. The huge sticker on the back doesn’t come off and looks super tacky. I would not purchase this again. I could have just bought a journal from the dollar store and it would be basically the same thing. It’s also really expensive for what it is.'\n",
       "\n",
       "'>>> Title: Not impressed at all... buy something else'\n",
       "\n",
       "'>>> Summary: Nothing special at all about this product'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print_summary(100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'>>> Review: Es una trilogia que se hace muy facil de leer. Me ha gustado, no me esperaba el final para nada'\n",
       "\n",
       "'>>> Title: Buena literatura para adolescentes'\n",
       "\n",
       "'>>> Summary: Muy facil de leer'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print_summary(0)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "Summarization (TensorFlow)",
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
